
import torch
from torch import nn

import cv2
from torchvision import transforms
import matplotlib.pyplot as plt
from ResNet import resnet18
# 可以用装饰器或者上下文的一个包
import contextlib
import time
import os
os.environ['KMP_DUPLICATE_LIB_OK'] = 'TRUE'



def yycc(path):

    print(path)
    # 检查是否有可用的 GPU，如果有则使用 GPU，否则使用 CPU
    if torch.cuda.is_available():
        device = torch.device("cuda")  # 使用 GPU
        print("CUDA is available. Using GPU.")
    else:
        device = torch.device("cpu")  # 使用 CPU
        print("CUDA is not available. Using CPU.")

    classes = {0: "half-ripe", 1: "raw", 2: "ripe", 3:"max_ripe"}


    # 假设你的模型定义在 `resnet18` 函数中并已保存最优模型为 './model/best.pth'
    model = resnet18().to(device)  # 根据你的模型定义调整 num_classes

    fc_inputs = model.fc.in_features

    model.fc = nn.Sequential(
        nn.Linear(fc_inputs, 128),  # 增加一个隐藏层
        nn.ReLU(),  # 添加 ReLU 激活函数
        # nn.Dropout(0.5),  # 添加 Dropout 防止过拟合
        nn.Linear(128, 4),  # 最终映射到 4 个输出类上
        nn.Softmax(dim=1)  # 分类任务使用 Softmax 输出概率分布
    ).to(device)

    model.load_state_dict(torch.load('./model/best.pth',weights_only=False))

    model.eval()  # 设置模型为评估模式
    # print('a')
    # 定义图像预处理步骤,torchvision.transforms，通常都需要PIL图像对象作为输入。
    preprocess = transforms.Compose([
        transforms.ToPILImage(),  # 先将opecv读入的numpy图像转换为PIL图像，pytorch识别图像需要PIL格式
        transforms.Resize((224, 224)),
        transforms.ToTensor(),  # 这一步会将图像的通道顺序从 (高度, 宽度, 通道) 转换为 (通道, 高度, 宽度)
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ])
    # print('b')
    # 加载并预处理图像
    # img_path = 'C:/Users/34346/Desktop/img0.jpg'  # 替换为你的图像路径
    # img_path = '../dataset_peach/valid/raw/img0.jpg'  # 替换为你的图像路径
    img_path = path  # 替换为你的图像路径
    image = cv2.imread(img_path)
    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)  # 将BGR格式转换为RGB格式
    # 预处理图像
    input_tensor = preprocess(image)  # 读入数据image的形状HWC，高度，宽度，通道；结果的形状是CHW, 通道，高度，宽度
    input_batch = input_tensor.unsqueeze(0).to(device)  # 创建一个批次 添加为BCHW

    # print('c')
    # 确保在不计算梯度的情况下进行推断
    with torch.no_grad():
        # output:Tensor(1, 3)
        output = model(input_batch)
        # probabilities:Tensor(3,)
        probabilities = torch.nn.functional.softmax(output[0], dim=0)

    # print('d')
    # 获取预测类别,返回最大概率值，最大概率值对应下标
    _, predicted_class = torch.max(probabilities, dim=0)
    # print('e')
    pre_class = classes[predicted_class.item()]

    # print(pre_class)

    return pre_class
    # plt.show()


# print(yycc('C:/Users/34346/Desktop/img0.jpg'))